import json

from openai import OpenAI
from conf.config import get_config
from function.functions import *


def main():
    """Query: 查看 app=grafana 且关键字包含 Error 的日志"""
    client = OpenAI(base_url="https://api.apiyi.com/v1")
    # 步骤一：把所有预定义的 function 传给 chatgpt
    user_input = input("输入查询指令：")
    message, functools = get_config()

    message.append({"role": "user", "content": user_input})

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=message,
        tools=functools,
        tool_choice="auto",
    )

    response_msg = response.choices[0].message
    tool_calls = response_msg.tool_calls
    print("\nChatGPT want to call function: ", tool_calls)

    # 步骤二：检查 LLM 是否调用了 function
    if tool_calls is None:
        print("not tool_calls")

    if tool_calls:
        available_functions = {
            "modify_config": modify_config,
            "restart_service": restart_service,
            "apply_manifest": apply_manifest
        }
        message.append(response_msg)

        # 步骤三：把每次 function 调用和返回的信息传给 model
        for tool_call in tool_calls:
            func_name = tool_call.function.name
            func_to_call = available_functions[func_name]
            func_args = json.loads(tool_call.function.arguments)
            print("func_args : ", func_args)
            func_response = func_to_call(**func_args)
            message.append(
                {
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": func_name,
                    "content": func_response
                }
            )

        print("msg : ", message)
        # 步骤四：把 function calling 的结果传给 model，进行对话
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=message
        )
        return response.choices[0].message.content


if __name__ == "__main__":
    print("LLM Resp: ", main())
